#!/usr/bin/env python3.10
import warnings
import argparse
from msvcrt import getwch
argParser = argparse.ArgumentParser()
argParser.add_argument("FILE", help="filename of the CAIE 9618 assembly program to be executed")
argParser.add_argument("--datafile", help="filename of a custom-format data file", default="", required=False)
args = argParser.parse_args()

# this input method is Windows-dependent: if you run Mac or Unix or Linux, seek for your Unicode getch equivalent.
def getch():
    return getwch()

def warnUser(msg: str):
    return warnings.warn(msg, UserWarning)

NUM_BITS = 32
BIT_CONVERT_MASK = int("0b"+"1"*NUM_BITS, base=0)
OPCODES = ["LDM","LDD","LDI","LDX","LDR","MOV","STO","STX","STI","ADD","SUB","INC","DEC","JMP","CMP","CMP",
           "CMI","JPE","JPN","IN","OUT","AND","XOR","OR","LSL","LSR","END","PRT"] # PRT is my own instruction: for debugging purposes

error = False

commands: list[str] = []
labels: dict[str, int] = dict()  # label_name --> line_num
symbs: dict[str, int] = dict()  # symbol_name --> mem_addr
mem: dict[int, int] = dict()  # mem_addr --> mem_val
lines: list[str] = []
OUTPUT = ""

def loadImmediate(s: str) -> int:
    if not s.startswith("#"):
        raise ValueError("Immediate Addressing must start with #")
    s = s[1:]
    if len(s) == 0:
        raise ValueError("Expected Number")
    if s[0] == "&":
        return int(s[1:], 16)
    if s[0] == "b":
        return int(s[1:], 2)
    return int(s)

def loadDirect(s: str) -> int:
    if s.isdecimal():
        x = int(s)
    else:
        if not (s in symbs.keys()):
            raise ValueError(f"Memory Label {s} INVALID")
        x = symbs[s]
    if not (x in mem.keys()):
        warnUser(f"Memory address {x} used before assignment. Default-initializing to 0.")
        mem[x] = 0
    return mem[x]

def loadIndirect(s: str) -> int: return loadDirect(str(loadDirect(s)))

def loadIndexed(s: str) -> int: return loadDirect(str(loadDirect(s) + ix))

def storeDirect(s: str, v: int):
    if s.isdecimal():
        x = int(s)
    else:
        if not (s in symbs.keys()):
            raise ValueError(f"Memory Label {s} invalid")
        x = symbs[s]
    mem[x] = v

def storeIndirect(s: str, v: int): storeDirect(str(loadDirect(s)), v)

def storeIndexed(s: str, v: int): storeDirect(str(loadDirect(s) + ix), v)

def loadImmediateOrDirect(s: str):
    if s.startswith("#"):
        return loadImmediate(s)
    return loadDirect(s)

def getAddress(s: str):
    if s.isdecimal():
        ret = int(oper) - 1 # our line count starts at 0.
    else:
        if not (oper in labels.keys()):
            raise ValueError(f"Invalid JMP label: {oper} on line {pc}")
        ret = labels[oper]
    return ret

# read from file
with open(args.FILE) as f:
    lines = f.readlines()
    while lines[-1] == "\n":
        lines.pop()

# read datafile
if args.datafile != "":
    with open(args.datafile) as f:
        dataLines = f.readlines()
        while dataLines[-1] == "\n":
            dataLines.pop()
    for line in dataLines:
        if line.startswith("//"):
            continue
        line, _, _ = line.partition("//") # discard comment
        while line.find("  ") != -1:
            line = line.replace("  ", " ")
        ls = line.strip().split()
        if len(ls) != 2:
            raise ValueError("Data Format: SYMBOL_NAME DATA_ADDRESS")
        symb, addr = ls[0], int(ls[1])
        if symb in symbs.keys():
            raise ValueError("Duplicate symbolic data address")
        symbs[symb] = addr

# pre-processing: separate label and opcode+operand
for i in range(len(lines)):
    line = lines[i].strip()
    line, _, _ = line.partition("//")
    if line.count(":") > 1:
        raise ValueError(f"Error on line {i+1}: More than one \":\" found.")
    label, _, command = line.rpartition(":")
    if label.isdecimal():
        raise ValueError(f"Error on line {i+1}: Labels cannot be numeric. Got: {label}")
    if label != "":
        if label in labels.keys():
            raise ValueError(f"Error on line {i+1}: Duplicate label {label}")
        labels[label] = i
    commands.append(command)

ix: int = 0
pc: int = 0
acc: int = 0
MAGIC_CMP = 114
cmp: int = MAGIC_CMP # Magic number... coz I'm lazy
MODE = "IN"

while True:
    command = commands[pc].strip()
    pc += 1
    parts = command.split()
    p = len(parts)
    acc &= BIT_CONVERT_MASK
    if p == 1:
        if command == "END":
            break
        elif command == "IN":
            if MODE == "OUT":
                MODE = "IN"
                print("\nIN: ", end = "")
            while True:
                ch = getch()
                if ord(ch) == 3 or ord(ch) == 27:
                    raise KeyboardInterrupt
                # print(ord(ch), ch)
                if ch.isascii() and 32 <= ord(ch) < 127:
                    break
                if ord(ch) == 13:
                    break
            print(ch, end = "", flush=True)
            acc = ord(ch)
        elif command == "OUT":
            if MODE == "IN":
                MODE = "OUT"
                print("\nOUT: ", end="")
            print(chr(acc), end = "", flush = True)
        elif command == "PRT":
            print("\nDEBUG: ", end="")
            if acc < (1<<(NUM_BITS-1)):
                print(acc, flush = True)
            elif acc < (1<<NUM_BITS):
                print(acc-(1<<NUM_BITS), flush = True)
            else:
                warnUser("Oh no something has gone very wrong!")
    elif p == 2:
        op, oper = parts
        op, oper = op.strip(), oper.strip()
        if op == "LDM":
            acc = loadImmediate(oper)
        elif op == "LDD":
            acc = loadDirect(oper)
        elif op == "LDI":
            acc = loadIndirect(oper)
        elif op == "LDX":
            ac = loadIndexed(oper)
        elif op == "MOV":
            if oper != "IX":
                raise ValueError(f"MOV operation must have operand \"IX\" on line {pc}")
            ix = acc
        elif op == "STO":
            storeDirect(oper, acc)
        elif op == "STX":
            storeIndexed(oper, acc)
        elif op == "STI":
            storeIndirect(oper, acc)
        elif op == "ADD":
            acc += loadImmediateOrDirect(oper)
        elif op == "SUB":
            acc -= loadImmediateOrDirect(oper)
        elif op == "INC":
            if oper == "ACC":
                acc += 1
            elif oper == "IX":
                ix += 1
            else:
                raise ValueError(f"INC operation must have operand ACC or IX on line {pc}")
        elif op == "DEC":
            if oper == "ACC":
                acc -= 1
            elif oper == "IX":
                ix -= 1
            else:
                raise ValueError(f"DEC operation must have operand ACC or IX on line {pc}")
        elif op == "JMP":
            pc = getAddress(oper)
        elif op == "CMP":
            val = loadImmediateOrDirect(oper)
            # Magic number --> no "cmp = acc - val"
            if acc > val:
                cmp = 1
            elif acc < val:
                cmp = -1
            else:
                cmp = 0
        elif op == "CMI":
            val = loadIndirect(oper)
            if acc > val:
                cmp = 1
            elif acc < val:
                cmp = -1
            else:
                cmp = 0
        elif op == "JPE":
            if cmp == MAGIC_CMP:
                raise ValueError(f"Conditional Jump before a compare instruction on line {pc}.")
            if cmp == 0:
                pc = getAddress(oper)
            else:
                pass
        elif op == "JPN":
            if cmp == MAGIC_CMP:
                raise ValueError(f"Conditional Jump before a compare instruction on line {pc}.")
            if cmp != 0:
                pc = getAddress(oper)
            else:
                pass
        elif op == "AND":
            val = loadImmediateOrDirect(oper)
            acc &= val
        elif op == "XOR":
            val = loadImmediateOrDirect(oper)
            acc ^= val
        elif op == "OR":
            val = loadImmediateOrDirect(oper)
            acc |= val
        elif op == "LSL":
            val = loadImmediate(oper)
            acc = acc << val
        elif op == "LSR": # Oh god. Python doesn't have logical shift operators...
            val = loadImmediate(oper)
            acc = acc >> val # acc is always positive, so logical shift == arithmetic shift

